# 读取图像和imu文件，转为rosbag (本例为4seasons数据集)
import rosbag
import rospy
import cv2
import time, sys, os
import roslib

roslib.load_manifest('sensor_msgs')
from sensor_msgs.msg import Image, Imu
from geometry_msgs.msg import Vector3
import numpy as np
from cv_bridge import CvBridge
from sensor_msgs.msg import CameraInfo


def quaternion_to_rotation_matrix(quat):
    print("q: ", quat)
    w = quat[0]
    x = quat[1]
    y = quat[2]
    z = quat[3]
    y2 = y * y
    x2 = x * x
    z2 = z * z
    xy = x * y
    xz = x * z
    yw = y * w
    zw = z * w
    yz = y * z
    xw = x * w

    rot_matrix = [[1.0 - 2 * y2 - 2 * z2, 2 * xy - 2 * zw, 2 * xz + 2 * yw],
                  [2 * xy + 2 * zw, 1.0 - 2 * x2 - 2 * z2, 2 * yz - 2 * xw],
                  [2 * xz - 2 * yw, 2 * yz + 2 * xw, 1.0 - 2 * x2 - 2 * y2]]
    return rot_matrix

def findFiles(root_dir, filter_type, reverse=False):
    """
    在指定目录查找指定类型文件 -> paths, names, files
    :param root_dir: 查找目录
    :param filter_type: 文件类型
    :param reverse: 是否返回倒序文件列表，默认为False
    :return: 路径、名称、文件全路径
    """

    separator = os.path.sep
    paths = []
    names = []
    files = []
    for parent, dirname, filenames in os.walk(root_dir):
        for filename in filenames:
            if filename.endswith(filter_type):
                paths.append(parent + separator)
                names.append(filename)
    for i in range(paths.__len__()):
        files.append(paths[i] + names[i])
    print(names.__len__().__str__() + " files have been found.")

    paths = np.array(paths)
    names = np.array(names)
    files = np.array(files)

    index = np.argsort(files)

    paths = paths[index]
    names = names[index]
    files = files[index]

    paths = list(paths)
    names = list(names)
    files = list(files)

    if reverse:
        paths.reverse()
        names.reverse()
        files.reverse()
    return paths, names, files


def write_imu(filename, topic):
    """return IMU data and timestamp of IMU"""
    file = open(filename, 'r')
    data = file.readlines()
    for i in range(len(data)):
        print("imu: ", i, '/', len(data))

        imu = Imu()
        angular_v = Vector3()
        linear_a = Vector3()

        line = data[i].rstrip('\n').split(' ')
        timestamp = int(line[0]) / 10e8
        # print(timestamp)
        imu_data = line[1:]
        assert (len(imu_data) == 6)
        angular_v.x = float(imu_data[0])
        angular_v.y = float(imu_data[1])
        angular_v.z = float(imu_data[2])
        linear_a.x = float(imu_data[3])
        linear_a.y = float(imu_data[4])
        linear_a.z = float(imu_data[5])
        imu_stamp = rospy.rostime.Time.from_sec(timestamp)  # according to the timestamp unit
        imu.header.stamp = imu_stamp
        imu.header.frame_id = "imu_link"
        imu.angular_velocity = angular_v
        imu.linear_acceleration = linear_a
        bag_out.write(topic, imu, imu_stamp)

    print("Total add %i imus!" % (len(data)))

def write_img(img_dir, img_type, topic, compress=True):
    paths, names, files = findFiles(img_dir, img_type)
    for i in range(len(files)):
        print("image: ", i, '/', len(files))

        frame_img = cv2.imread(files[i], cv2.IMREAD_GRAYSCALE)
        timestamp = int(names[i].split(".")[0]) / 10e8
        # print(timestamp)

        ros_ts = rospy.rostime.Time.from_sec(timestamp)
        if compress:
            ros_img = cb.cv2_to_compressed_imgmsg(frame_img, dst_format='png')
        else:
            ros_img = cb.cv2_to_imgmsg(frame_img, encoding='mono8')
        ros_img.header.stamp = ros_ts
        bag_out.write(topic, ros_img, ros_ts)


if __name__ == '__main__':
    source_dir = "/media/daybeha/软件/win/桌面/dataset/4Seasons/recording_2021-02-25_13-51-57"
    img0_dir = source_dir + "/distorted_images/cam0"  # 图像所在文件夹路径
    img1_dir = source_dir + "/distorted_images/cam1"  # 图像所在文件夹路径
    img_type = "png"  # 影像类型
    img0_topic = "/image_left/image_raw"  # Topic名称
    img1_topic = "/image_right/image_raw"  # Topic名称
    compress = True
    compress = False

    imu_dir = source_dir + "/imu.txt"
    imu_topic = "/imu"

    bag_path = "office_loop_6_train.bag"  # 输出Bag路径


    bag_out = rosbag.Bag(bag_path, 'w')
    if compress:
        img0_topic += "/compressed"
        img1_topic += "/compressed"

    cb = CvBridge()

    write_imu(imu_dir, imu_topic)
    write_img(img0_dir, img_type, img0_topic, compress)
    write_img(img1_dir, img_type, img1_topic, compress)

    bag_out.close()
